#!/usr/bin/python
import collections
import numpy as np
import rospy
from scipy.signal import argrelextrema
from math import degrees, sqrt, pi
from std_msgs.msg import Float32

from sensor_msgs.msg import Imu
from math import factorial
from math import atan2

class tack_params():
    def __init__(self):
        self.wave_period_pub = rospy.Publisher('wave_period', Float32, queue_size=10)

        self.roll_angle_pub = rospy.Publisher('roll_angle', Float32, queue_size=10)

        rospy.init_node('wave_period_node', anonymous=True)

        rospy.Subscriber('/imu/data', Imu, self.update_AccZ)
        self.AccZ = collections.deque(maxlen =500)
        self.Acc_Y = 0
        self.Acc_Z = 0
        self.Acc_X = 0 # trying to repair the code

        self.rate = rospy.Rate(10)

        self.period_publisher()

    def period_publisher(self):
        """
        Publish wave period and roll angle
        """

        while not rospy.is_shutdown():
            if len(self.AccZ) <500:
                self.rate.sleep()
                continue

            signal = np.array(self.AccZ)
            smooth = savitzky_golay(signal, 51, 3)
            index_max = argrelextrema(smooth, np.greater)[0]
            index_min = argrelextrema(smooth, np.less)[0]

            wave_period_max = (index_max[-1]-index_max[0])/(len(index_max)-1)
            wave_period_min = (index_min[-1]-index_min[0])/(len(index_min)-1)
            wave_period = (wave_period_max + wave_period_min)/2.0
            #in seconds, divide by imu frequency

            wave_period = wave_period/100.0

            #time since last cress

            t_cress = (len(signal) - index_min[-1]) / 100.0
            t_crest = (index_min[-1] - index_min [-2]) / 100.0

            #position on wave relative to wave period

            position = t_cress/wave_period
            roll_angle = atan2(self.Acc_Y, sqrt(self.Acc_Z**2 + self.Acc_X**2)) # reparing the code

            self.wave_period_pub.publish(position)
            self.roll_angle_pub.publish(degrees(roll_angle))

            self.rate.sleep()

    def update_AccZ(self, msg):
        self.AccZ.append(msg.linear_acceleration.z)
        self.Acc_X = msg.linear_acceleration.x # trying to repair the code
        self.Acc_Y = msg.linear_acceleration.y
        self.Acc_Z = msg.linear_acceleration.z


def savitzky_golay(y, window_size, order, deriv=0, rate=1):
    try:
        window_size = np.abs(np.int(window_size))
        order = np.abs(np.int(order))
    except ValueError:
        raise ValueError("window_size and order have to be of type int")
    if window_size % 2 != 1 or window_size < 1:
        raise TypeError("window_size size must be a positive odd number")
    if window_size < order + 2:
        raise TypeError("window_size is too small for the polynomials order")

    order_range = range(order+1)
    half_window = (window_size -1) // 2
    # precompute coefficients
    b = np.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)])
    m = np.linalg.pinv(b).A[deriv] * rate**deriv * factorial(deriv)
    # pad the signal at the extremes with
    # values taken from the signal itself
    firstvals = y[0] - np.abs( y[1:half_window+1][::-1] - y[0] )
    lastvals = y[-1] + np.abs(y[-half_window-1:-1][::-1] - y[-1])
    y = np.concatenate((firstvals, y, lastvals))
    return np.convolve( m[::-1], y, mode='valid')

if __name__ == '__main__':
    try:
        tack_params()
    except rospy.ROSInterruptException:
        pass
